/* benchmark guru interfaces */
#include <Configuration.h>
#include <mpi/communicator.h>
#include <mpi/collectives.h>
#include <Message/OpenMP.h>
#include <Utilities/Timer.h>
#include <OhmmsPETE/OhmmsMatrix.h>
#include <fft/fft.h>
#include <bitset>

namespace APPNAMESPACE {

/** class to perform FFT benchmark
 * 
 * PREC precision of the data, double, float
 * T datatype: double, complex<double>, float, complex<float>
 * ffteng FFT engine as defined in FFT/fft.h
 * Using Matrix<T>
 */
template<typename PREC, typename T,bool domany, unsigned ffteng>
struct FFT_PFOR
{
  ///save the enum
  enum {MANY_FFT=domany};
  ///data to be transformed
  Matrix<T> target;
  ///FFT engine which can do multiple ffts of domany==true
  FFT_DFT<PREC,T,domany,ffteng> fft_engine;
  ///number of ffts
  int num_ffts;
  ///number of ffts
  int num_blocks;
  ///memory stride to perform blocked FFTs
  int stride_size;

  /** constructor
   * @param fft_dim FFT dimension
   * @param howmany number of ffts
   * @param buffersize fft_dim+2 for T=real
   * @param nb nunber of simultaneous FFTs 
   */
  inline FFT_PFOR(int fft_dim, int howmany, int buffersize, int nb=1)
    :target(howmany*nb,buffersize) 
     , fft_engine(fft_dim,target.data(),howmany)
     , num_ffts(howmany)
     , num_blocks(nb)
     , stride_size(buffersize*howmany)
  {
  }

  inline void benchmark(int niters)
  {
    for(int iter=0; iter<niters; ++iter)
    {
#pragma omp for nowait
      for(int k=0; k<num_blocks; ++k)
      {
        fft_engine.transform_forward(target.data()+stride_size*k,num_ffts);
      }

#pragma omp for nowait
      for(int k=0; k<num_blocks; ++k)
      {
        fft_engine.transform_backward(target.data()+stride_size*k,num_ffts);
      }
    }
  }
};
}

using namespace APPNAMESPACE;
int main(int argc, char** argv)
{
  mpi::environment env(argc,argv);
  mpi::communicator mycomm;
  OhmmsInfo ohmms("fft_bench",mycomm.rank(),0,1);

  ///1D FFT dimension
  int fft_dim=12288;
  ///total number of ffts
  int fft_tot=512;
  ///number of fft blocks per call
  int fft_block=512;
  int niters=10;
  unsigned plan_mode = FFTW_MEASURE;

  bitset<MAX_FFT_ENG> test_it(0);
  test_it.set(FFTW_ENG,1);

  //cout.setf(std::ios::scientific, std::ios::floatfield);
  int ic=0;
  while(ic<argc)
  {
    string a(argv[ic]);
    if(a.find("-s")<a.size()) //get fft dimension
      fft_dim=atoi(argv[++ic]);
    else if(a.find("-n")<a.size()) //get how many fft to perform
      fft_tot=atoi(argv[++ic]);
    else if(a.find("-b")<a.size()) //number of simultaneous trans.
      fft_block=atoi(argv[++ic]);
    else if(a.find("-i")<a.size()) 
      niters=atoi(argv[++ic]);
    else if(a.find("disable-fftw")<a.size())
      test_it.set(FFTW_ENG,0);
    else if(a.find("enable-essl")<a.size())
      test_it.set(FFTESSL_ENG,1);
    else if(a.find("enable-mkl")<a.size())
      test_it.set(FFTMKL_ENG,1);
    ++ic;
  }

  fft_block=std::min(fft_block,fft_tot);
  fft_tot=(fft_tot/fft_block)*fft_block;
  int num_fft_blocks=fft_tot/fft_block;

  double gflops_c=(log(static_cast<double>(fft_dim))/log(2.))*5.*fft_dim*fft_tot*1e-9;
  double gflops_r=gflops_c*0.5;

  app_log() << "#  Testing FFT dim= " << fft_dim 
    << " number of fft = " << fft_tot
    << " fft bundles = " << fft_block
    << " fft blocks = " << num_fft_blocks
    << " iterations= " << niters 
    << " MFLOPS(complex)= " << static_cast<int>(gflops_c*1000)
    << endl;
  app_log() << "# MPI " << mycomm.size() << " OMP_NUM_THREADS = " << omp_get_max_threads() << endl;


  typedef double real_type;
  typedef std::complex<double> complex_type;
  FFT_PFOR<real_type,complex_type,true,FFTW_ENG> many(fft_dim,fft_block,fft_dim+2,num_fft_blocks);

#pragma omp parallel
  {
    many.benchmark(niters);
  }

  return 0;
}

